package com.devbobcorn.nekoration.client.rendering.chunks;

import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.Set;
import java.util.concurrent.CancellationException;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
import java.util.stream.Collectors;

import javax.annotation.Nullable;

import com.google.common.collect.Lists;
import com.google.common.collect.Sets;
import com.google.common.primitives.Doubles;
import com.mojang.blaze3d.matrix.MatrixStack;

import net.minecraft.block.BlockRenderType;
import net.minecraft.block.BlockState;
import net.minecraft.client.Minecraft;
import net.minecraft.client.renderer.ActiveRenderInfo;
import net.minecraft.client.renderer.BlockModelRenderer;
import net.minecraft.client.renderer.BlockRendererDispatcher;
import net.minecraft.client.renderer.BufferBuilder;
import net.minecraft.client.renderer.RegionRenderCacheBuilder;
import net.minecraft.client.renderer.RenderType;
import net.minecraft.client.renderer.RenderTypeLookup;
import net.minecraft.client.renderer.WorldRenderer;
import net.minecraft.client.renderer.chunk.ChunkRenderCache;
import net.minecraft.client.renderer.chunk.ChunkRenderDispatcher;
import net.minecraft.client.renderer.chunk.VisGraph;
import net.minecraft.client.renderer.tileentity.TileEntityRenderer;
import net.minecraft.client.renderer.tileentity.TileEntityRendererDispatcher;
import net.minecraft.client.renderer.vertex.DefaultVertexFormats;
import net.minecraft.client.renderer.vertex.VertexBuffer;
import net.minecraft.crash.CrashReport;
import net.minecraft.fluid.FluidState;
import net.minecraft.tileentity.TileEntity;
import net.minecraft.util.Direction;
import net.minecraft.util.Util;
import net.minecraft.util.math.AxisAlignedBB;
import net.minecraft.util.math.BlockPos;
import net.minecraft.util.math.vector.Vector3d;
import net.minecraft.world.World;
import net.minecraft.world.chunk.Chunk;
import net.minecraft.world.chunk.ChunkStatus;
import net.minecraftforge.api.distmarker.Dist;
import net.minecraftforge.api.distmarker.OnlyIn;

// This class is a modified version of vanilla ChunkRenderDispatcher.ChunkRender,
// which is indepent of ChunkRenderDispatcher and can be easily manipulated.
@OnlyIn(Dist.CLIENT)
public class ChunkModelRender implements net.minecraftforge.client.extensions.IForgeRenderChunk {
    public final AtomicReference<CompiledModelChunk> compiled = new AtomicReference<>(
            CompiledModelChunk.UNCOMPILED);
    @Nullable
    private ChunkModelRender.RebuildTask lastRebuildTask;
    @Nullable
    private ChunkModelRender.SortTransparencyTask lastResortTransparencyTask;
    private final Set<TileEntity> globalBlockEntities = Sets.newHashSet();
    private final Map<RenderType, VertexBuffer> buffers = RenderType.chunkBufferLayers().stream()
            .collect(Collectors.toMap((p_228934_0_) -> {
                return p_228934_0_;
            }, (p_228933_0_) -> {
                return new VertexBuffer(DefaultVertexFormats.BLOCK);
            }));
    public AxisAlignedBB bb;
    private int lastFrame = -1;
    private boolean dirty = true;
    private final BlockPos.Mutable origin = new BlockPos.Mutable(-1, -1, -1);
    private final BlockPos.Mutable[] relativeOrigins = Util.make(new BlockPos.Mutable[6], (p_228932_0_) -> {
        for (int i = 0; i < p_228932_0_.length; ++i) {
            p_228932_0_[i] = new BlockPos.Mutable();
        }
    });
    private boolean playerChanged;

    // Additional fields...
    private World level;
    private final WorldRenderer renderer;
    private Vector3d camera = Vector3d.ZERO;
    public final RegionRenderCacheBuilder fixedBuffers;

    // Constructor with additional data
    public ChunkModelRender(World wd, WorldRenderer rd){
        this.level = wd;
        this.renderer = rd;
        this.fixedBuffers = new RegionRenderCacheBuilder();
    }

    public void setCamera(Vector3d cam){
        this.camera = cam;
    }

    private CompletableFuture<Void> doUploadChunkLayer(BufferBuilder p_228904_1_, VertexBuffer p_228904_2_) {
        return p_228904_2_.uploadLater(p_228904_1_);
    }

    public void compileModel(){
        this.compileSync();
    }
    // Modification end

    private boolean doesChunkExistAt(BlockPos p_228930_1_) {
        return this.level.getChunk(p_228930_1_.getX() >> 4, p_228930_1_.getZ() >> 4,
                ChunkStatus.FULL, false) != null;
    }

    public boolean hasAllNeighbors() {
        if (!(this.getDistToPlayerSqr() > 576.0D)) {
            return true;
        } else {
            return this.doesChunkExistAt(this.relativeOrigins[Direction.WEST.ordinal()])
                    && this.doesChunkExistAt(this.relativeOrigins[Direction.NORTH.ordinal()])
                    && this.doesChunkExistAt(this.relativeOrigins[Direction.EAST.ordinal()])
                    && this.doesChunkExistAt(this.relativeOrigins[Direction.SOUTH.ordinal()]);
        }
    }

    public boolean setFrame(int p_178577_1_) {
        if (this.lastFrame == p_178577_1_) {
            return false;
        } else {
            this.lastFrame = p_178577_1_;
            return true;
        }
    }

    public VertexBuffer getBuffer(RenderType p_228924_1_) {
        return this.buffers.get(p_228924_1_);
    }

    public void setOrigin(int x, int y, int z) {
        if (x != this.origin.getX() || y != this.origin.getY()
                || z != this.origin.getZ()) {
            this.reset();
            this.origin.set(x, y, z);
            this.bb = new AxisAlignedBB((double) x, (double) y, (double) z, (double) (x + 16), (double) (y + 16), (double) (z + 16));

            for (Direction direction : Direction.values()) {
                this.relativeOrigins[direction.ordinal()].set(this.origin).move(direction, 16);
            }

        }
    }

    @SuppressWarnings("resource")
    protected double getDistToPlayerSqr() {
        ActiveRenderInfo activerenderinfo = Minecraft.getInstance().gameRenderer.getMainCamera();
        double d0 = this.bb.minX + 8.0D - activerenderinfo.getPosition().x;
        double d1 = this.bb.minY + 8.0D - activerenderinfo.getPosition().y;
        double d2 = this.bb.minZ + 8.0D - activerenderinfo.getPosition().z;
        return d0 * d0 + d1 * d1 + d2 * d2;
    }

    private void beginLayer(BufferBuilder buffer) {
        buffer.begin(7, DefaultVertexFormats.BLOCK);
    }

    public CompiledModelChunk getCompiledChunk() {
        return this.compiled.get();
    }

    private void reset() {
        this.cancelTasks();
        this.compiled.set(CompiledModelChunk.UNCOMPILED);
        this.dirty = true;
    }

    public void releaseBuffers() {
        this.reset();
        this.buffers.values().forEach(VertexBuffer::close);
    }

    public BlockPos getOrigin() {
        return this.origin;
    }

    public void setDirty(boolean p_178575_1_) {
        boolean flag = this.dirty;
        this.dirty = true;
        this.playerChanged = p_178575_1_ | (flag && this.playerChanged);
    }

    public void setNotDirty() {
        this.dirty = false;
        this.playerChanged = false;
    }

    public boolean isDirty() {
        return this.dirty;
    }

    public boolean isDirtyFromPlayer() {
        return this.dirty && this.playerChanged;
    }

    public BlockPos getRelativeOrigin(Direction p_181701_1_) {
        return this.relativeOrigins[p_181701_1_.ordinal()];
    }

    public boolean resortTransparency(RenderType p_228925_1_) {
        CompiledModelChunk compiledchunk = this.getCompiledChunk();
        if (this.lastResortTransparencyTask != null) {
            this.lastResortTransparencyTask.cancel();
        }

        if (!compiledchunk.hasLayer.contains(p_228925_1_)) {
            return false;
        } else {
            this.lastResortTransparencyTask = new ChunkModelRender.SortTransparencyTask(
                    new net.minecraft.util.math.ChunkPos(getOrigin()), this.getDistToPlayerSqr(),
                    compiledchunk);
            //p_228925_2_.schedule(this.lastResortTransparencyTask);
            this.lastResortTransparencyTask.doTask(this.fixedBuffers); // Just do it straight away
            return true;
        }
    }

    protected void cancelTasks() {
        if (this.lastRebuildTask != null) {
            this.lastRebuildTask.cancel();
            this.lastRebuildTask = null;
        }

        if (this.lastResortTransparencyTask != null) {
            this.lastResortTransparencyTask.cancel();
            this.lastResortTransparencyTask = null;
        }
    }

    public ChunkModelRender.ChunkRenderTask createCompileTask() {
        this.cancelTasks();
        BlockPos blockpos = this.origin.immutable();
        ChunkRenderCache chunkrendercache = createRegionRenderCache(this.level, blockpos.offset(-1, -1, -1), blockpos.offset(16, 16, 16), 1);
        this.lastRebuildTask = new ChunkModelRender.RebuildTask(new net.minecraft.util.math.ChunkPos(getOrigin()), this.getDistToPlayerSqr(), chunkrendercache);
        return this.lastRebuildTask;
    }

    public void rebuildChunkSync(ChunkRenderDispatcher p_228929_1_) {
        ChunkModelRender.ChunkRenderTask chunkrendertask = this.createCompileTask();
        //p_228929_1_.schedule(chunkrendertask);
        chunkrendertask.doTask(this.fixedBuffers);
    }

    private void updateGlobalBlockEntities(Set<TileEntity> p_228931_1_) {
        Set<TileEntity> set = Sets.newHashSet(p_228931_1_);
        Set<TileEntity> set1 = Sets.newHashSet(this.globalBlockEntities);
        set.removeAll(this.globalBlockEntities);
        set1.removeAll(p_228931_1_);
        this.globalBlockEntities.clear();
        this.globalBlockEntities.addAll(p_228931_1_);
        ChunkModelRender.this.renderer.updateGlobalBlockEntities(set1, set);
    }

    public void compileSync() {
        ChunkModelRender.ChunkRenderTask chunkrendertask = this.createCompileTask();
        chunkrendertask.doTask(this.fixedBuffers);
    }

    @OnlyIn(Dist.CLIENT)
    public abstract class ChunkRenderTask implements Comparable<ChunkModelRender.ChunkRenderTask> {
        protected final double distAtCreation;
        protected final AtomicBoolean isCancelled = new AtomicBoolean(false);
        protected java.util.Map<net.minecraft.util.math.BlockPos, net.minecraftforge.client.model.data.IModelData> modelData;

        public ChunkRenderTask(double p_i226023_2_) {
            this(null, p_i226023_2_);
        }

        @SuppressWarnings("resource")
        public ChunkRenderTask(@Nullable net.minecraft.util.math.ChunkPos pos, double p_i226023_2_) {
            this.distAtCreation = p_i226023_2_;
            if (pos == null) {
                this.modelData = java.util.Collections.emptyMap();
            } else {
                this.modelData = net.minecraftforge.client.model.ModelDataManager.getModelData(Minecraft.getInstance().level, pos);
            }
        }

        public abstract CompletableFuture<ModelChunkTaskResult> doTask(RegionRenderCacheBuilder p_225618_1_);

        public abstract void cancel();

        public int compareTo(ChunkModelRender.ChunkRenderTask p_compareTo_1_) {
            return Doubles.compare(this.distAtCreation, p_compareTo_1_.distAtCreation);
        }

        public net.minecraftforge.client.model.data.IModelData getModelData(net.minecraft.util.math.BlockPos pos) {
            return modelData.getOrDefault(pos, net.minecraftforge.client.model.data.EmptyModelData.INSTANCE);
        }
    }

    @OnlyIn(Dist.CLIENT)
    class RebuildTask extends ChunkModelRender.ChunkRenderTask {
        @Nullable
        protected ChunkRenderCache region;

        @Deprecated
        public RebuildTask(double p_i226024_2_, @Nullable ChunkRenderCache p_i226024_4_) {
            this(null, p_i226024_2_, p_i226024_4_);
        }

        public RebuildTask(@Nullable net.minecraft.util.math.ChunkPos pos, double p_i226024_2_,
                @Nullable ChunkRenderCache p_i226024_4_) {
            super(pos, p_i226024_2_);
            this.region = p_i226024_4_;
        }

        public CompletableFuture<ModelChunkTaskResult> doTask(RegionRenderCacheBuilder p_225618_1_) {
            if (this.isCancelled.get()) {
                return CompletableFuture.completedFuture(ModelChunkTaskResult.CANCELLED);
            } else if (!ChunkModelRender.this.hasAllNeighbors()) {
                this.region = null;
                ChunkModelRender.this.setDirty(false);
                this.isCancelled.set(true);
                return CompletableFuture.completedFuture(ModelChunkTaskResult.CANCELLED);
            } else if (this.isCancelled.get()) {
                return CompletableFuture.completedFuture(ModelChunkTaskResult.CANCELLED);
            } else {
                Vector3d cam = ChunkModelRender.this.camera;
                float camX = (float) cam.x;
                float camY = (float) cam.y;
                float camZ = (float) cam.z;
                CompiledModelChunk compiledchunk = new CompiledModelChunk();
                Set<TileEntity> set = this.compile(camX, camY, camZ, compiledchunk, p_225618_1_);
                ChunkModelRender.this.updateGlobalBlockEntities(set);
                if (this.isCancelled.get()) {
                    return CompletableFuture.completedFuture(ModelChunkTaskResult.CANCELLED);
                } else {
                    List<CompletableFuture<Void>> list = Lists.newArrayList();
                    compiledchunk.hasLayer.forEach((p_228943_3_) -> {
                        list.add(ChunkModelRender.this.doUploadChunkLayer(p_225618_1_.builder(p_228943_3_),
                                ChunkModelRender.this.getBuffer(p_228943_3_)));
                    });
                    return Util.sequence(list).handle((p_228941_2_, p_228941_3_) -> {
                        if (p_228941_3_ != null && !(p_228941_3_ instanceof CancellationException)
                                && !(p_228941_3_ instanceof InterruptedException)) {
                            Minecraft.getInstance()
                                    .delayCrash(CrashReport.forThrowable(p_228941_3_, "Rendering chunk"));
                        }

                        if (this.isCancelled.get()) {
                            return ModelChunkTaskResult.CANCELLED;
                        } else {
                            ChunkModelRender.this.compiled.set(compiledchunk);
                            return ModelChunkTaskResult.SUCCESSFUL;
                        }
                    });
                }
            }
        }

        private Set<TileEntity> compile(float p_228940_1_, float p_228940_2_, float p_228940_3_,
                CompiledModelChunk p_228940_4_, RegionRenderCacheBuilder p_228940_5_) {
            BlockPos blockpos = ChunkModelRender.this.origin.immutable();
            BlockPos blockpos1 = blockpos.offset(15, 15, 15);
            VisGraph visgraph = new VisGraph();
            Set<TileEntity> set = Sets.newHashSet();
            ChunkRenderCache chunkrendercache = this.region;
            this.region = null;
            MatrixStack matrixstack = new MatrixStack();
            if (chunkrendercache != null) {
                BlockModelRenderer.enableCaching();
                Random random = new Random();
                BlockRendererDispatcher blockrendererdispatcher = Minecraft.getInstance().getBlockRenderer();

                for (BlockPos blockpos2 : BlockPos.betweenClosed(blockpos, blockpos1)) {
                    BlockState blockstate = chunkrendercache.getBlockState(blockpos2);
                    if (blockstate.isSolidRender(chunkrendercache, blockpos2)) {
                        visgraph.setOpaque(blockpos2);
                    }

                    if (blockstate.hasTileEntity()) {
                        TileEntity tileentity = chunkrendercache.getBlockEntity(blockpos2, Chunk.CreateEntityType.CHECK);
                        if (tileentity != null) {
                            this.handleBlockEntity(p_228940_4_, set, tileentity);
                        }
                    }

                    FluidState fluidstate = chunkrendercache.getFluidState(blockpos2);
                    net.minecraftforge.client.model.data.IModelData modelData = getModelData(blockpos2);
                    for (RenderType rendertype : RenderType.chunkBufferLayers()) {
                        net.minecraftforge.client.ForgeHooksClient.setRenderLayer(rendertype);
                        if (!fluidstate.isEmpty() && RenderTypeLookup.canRenderInLayer(fluidstate, rendertype)) {
                            BufferBuilder bufferbuilder = p_228940_5_.builder(rendertype);
                            if (p_228940_4_.hasLayer.add(rendertype)) {
                                ChunkModelRender.this.beginLayer(bufferbuilder);
                            }

                            if (blockrendererdispatcher.renderLiquid(blockpos2, chunkrendercache, bufferbuilder,
                                    fluidstate)) {
                                p_228940_4_.isCompletelyEmpty = false;
                                p_228940_4_.hasBlocks.add(rendertype);
                            }
                        }

                        if (blockstate.getRenderShape() != BlockRenderType.INVISIBLE
                                && RenderTypeLookup.canRenderInLayer(blockstate, rendertype)) {
                            RenderType rendertype1 = rendertype;
                            BufferBuilder bufferbuilder2 = p_228940_5_.builder(rendertype1);
                            if (p_228940_4_.hasLayer.add(rendertype1)) {
                                ChunkModelRender.this.beginLayer(bufferbuilder2);
                            }

                            matrixstack.pushPose();
                            matrixstack.translate((double) (blockpos2.getX() & 15), (double) (blockpos2.getY() & 15),
                                    (double) (blockpos2.getZ() & 15));
                            if (blockrendererdispatcher.renderModel(blockstate, blockpos2, chunkrendercache,
                                    matrixstack, bufferbuilder2, true, random, modelData)) {
                                p_228940_4_.isCompletelyEmpty = false;
                                p_228940_4_.hasBlocks.add(rendertype1);
                            }

                            matrixstack.popPose();
                        }
                    }
                }
                net.minecraftforge.client.ForgeHooksClient.setRenderLayer(null);

                if (p_228940_4_.hasBlocks.contains(RenderType.translucent())) {
                    BufferBuilder bufferbuilder1 = p_228940_5_.builder(RenderType.translucent());
                    bufferbuilder1.sortQuads(p_228940_1_ - (float) blockpos.getX(),
                            p_228940_2_ - (float) blockpos.getY(), p_228940_3_ - (float) blockpos.getZ());
                    p_228940_4_.transparencyState = bufferbuilder1.getState();
                }

                p_228940_4_.hasLayer.stream().map(p_228940_5_::builder).forEach(BufferBuilder::end);
                BlockModelRenderer.clearCache();
            }

            p_228940_4_.visibilitySet = visgraph.resolve();
            return set;
        }

        private <E extends TileEntity> void handleBlockEntity(CompiledModelChunk p_228942_1_, Set<TileEntity> p_228942_2_, E p_228942_3_) {
            TileEntityRenderer<E> tileentityrenderer = TileEntityRendererDispatcher.instance.getRenderer(p_228942_3_);
            if (tileentityrenderer != null) {
                if (tileentityrenderer.shouldRenderOffScreen(p_228942_3_)) {
                    p_228942_2_.add(p_228942_3_);
                } else
                    p_228942_1_.renderableBlockEntities.add(p_228942_3_); // FORGE: Fix MC-112730
            }

        }

        public void cancel() {
            this.region = null;
            if (this.isCancelled.compareAndSet(false, true)) {
                ChunkModelRender.this.setDirty(false);
            }

        }
    }

    @OnlyIn(Dist.CLIENT)
    class SortTransparencyTask extends ChunkModelRender.ChunkRenderTask {
        private final CompiledModelChunk compiledChunk;

        @Deprecated
        public SortTransparencyTask(double p_i226025_2_, CompiledModelChunk p_i226025_4_) {
            this(null, p_i226025_2_, p_i226025_4_);
        }

        public SortTransparencyTask(@Nullable net.minecraft.util.math.ChunkPos pos, double p_i226025_2_, CompiledModelChunk p_i226025_4_) {
            super(pos, p_i226025_2_);
            this.compiledChunk = p_i226025_4_;
        }

        public CompletableFuture<ModelChunkTaskResult> doTask(RegionRenderCacheBuilder p_225618_1_) {
            if (this.isCancelled.get()) {
                return CompletableFuture.completedFuture(ModelChunkTaskResult.CANCELLED);
            } else if (!ChunkModelRender.this.hasAllNeighbors()) {
                this.isCancelled.set(true);
                return CompletableFuture.completedFuture(ModelChunkTaskResult.CANCELLED);
            } else if (this.isCancelled.get()) {
                return CompletableFuture.completedFuture(ModelChunkTaskResult.CANCELLED);
            } else {
                Vector3d vector3d = ChunkModelRender.this.camera;
                float f = (float) vector3d.x;
                float f1 = (float) vector3d.y;
                float f2 = (float) vector3d.z;
                BufferBuilder.State bufferbuilder$state = this.compiledChunk.transparencyState;
                if (bufferbuilder$state != null && this.compiledChunk.hasBlocks.contains(RenderType.translucent())) {
                    BufferBuilder bufferbuilder = p_225618_1_.builder(RenderType.translucent());
                    ChunkModelRender.this.beginLayer(bufferbuilder);
                    bufferbuilder.restoreState(bufferbuilder$state);
                    bufferbuilder.sortQuads(f - (float) ChunkModelRender.this.origin.getX(),
                            f1 - (float) ChunkModelRender.this.origin.getY(), f2 - (float) ChunkModelRender.this.origin.getZ());
                    this.compiledChunk.transparencyState = bufferbuilder.getState();
                    bufferbuilder.end();
                    if (this.isCancelled.get()) {
                        return CompletableFuture.completedFuture(ModelChunkTaskResult.CANCELLED);
                    } else {
                        CompletableFuture<ModelChunkTaskResult> completablefuture = ChunkModelRender.this
                                .doUploadChunkLayer(p_225618_1_.builder(RenderType.translucent()),
                                        ChunkModelRender.this.getBuffer(RenderType.translucent()))
                                .thenApply((p_228947_0_) -> {
                                    return ModelChunkTaskResult.CANCELLED;
                                });
                        return completablefuture.handle((p_228946_1_, p_228946_2_) -> {
                            if (p_228946_2_ != null && !(p_228946_2_ instanceof CancellationException)
                                    && !(p_228946_2_ instanceof InterruptedException)) {
                                Minecraft.getInstance()
                                        .delayCrash(CrashReport.forThrowable(p_228946_2_, "Rendering chunk"));
                            }

                            return this.isCancelled.get() ? ModelChunkTaskResult.CANCELLED
                                    : ModelChunkTaskResult.SUCCESSFUL;
                        });
                    }
                } else {
                    return CompletableFuture.completedFuture(ModelChunkTaskResult.CANCELLED);
                }
            }
        }

        public void cancel() {
            this.isCancelled.set(true);
        }
    }
}

